#include "nnslitememory.h"

#define NNSLITE_MEMORY_ALLOCRECORD_NODESIZE 4096

static bool gMemoryContextsInitialized = false;
static nnsliteMemoryContext gSystemMemContext;
static nnsliteMemoryContext gDefaultMemContext;

typedef struct AllocRecord
{
	void* addr;
	u32   size;
	u64   id;
}AllocRecord;

static NNSLITE_RET CreateSystemMemoryContext();
static NNSLITE_RET CreateDefaultMemoryContext();
static void* DefaultAllocMem(size_t size);
static void  DefaultFreeMem(void *p);
static NNSLITE_RET DefaultCheckMem();

NNSLITE_RET nnsliteMemoryInit(void)
{
	NNSLITE_RET ret;
	/* Create the system memory context 
	 * Note: The system memory context must be created at first ! */
	ret = CreateSystemMemoryContext();
	/* Create the default memory context */
	ret = CreateDefaultMemoryContext();
	NNSLITE_ON_FAILED_RETURN(ret);
	/* Set the flag */
	gMemoryContextsInitialized = true;
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMemoryExit(void)
{
	NNSLITE_RET ret;
	/* Free the default memory context */
	ret = nnsliteBlistFree(&gDefaultMemContext.allocList);
	NNSLITE_ON_FAILED_RETURN(ret);
	return NNSLITE_RET_OK;
}

nnsliteMemoryContext* nnsliteMemoryGetDefaultContext()
{
#ifdef _DEBUG
	return &gDefaultMemContext;
#else
	return &gSystemMemContext;
#endif
}

static NNSLITE_RET CreateSystemMemoryContext()
{
	gSystemMemContext.allocMem = malloc;
	gSystemMemContext.freeMem = free;
	return NNSLITE_RET_OK;
}

static NNSLITE_RET CreateDefaultMemoryContext()
{
	NNSLITE_RET ret;
	ret = nnsliteBlistCreate(&gDefaultMemContext.allocList, sizeof(AllocRecord), NNSLITE_MEMORY_ALLOCRECORD_NODESIZE, &gSystemMemContext);
	NNSLITE_ON_FAILED_RETURN(ret);
	gDefaultMemContext.idcounter = 0;
	gDefaultMemContext.allocMem = DefaultAllocMem;
	gDefaultMemContext.freeMem = DefaultFreeMem;
	gDefaultMemContext.checkMem = DefaultCheckMem;
	return NNSLITE_RET_OK;
}

static void* DefaultAllocMem(size_t size)
{
	nnsliteBlist* blist;
	AllocRecord record;
	void* pnew;
	/* Allocate the memory by the system memory context */
	pnew = gSystemMemContext.allocMem(size);
	/* Record this allocation */
	blist = &gDefaultMemContext.allocList;
	record.addr = pnew;
	record.size = (u32)size;
	record.id = gDefaultMemContext.idcounter++;
	nnsliteBlistInsert(blist, &record);
	return pnew;
}

static void DefaultFreeMem(void *p)
{
	nnsliteBlistIterator iter;
	AllocRecord record;
	nnsliteBlist* blist = &gDefaultMemContext.allocList;
	/* Find which allocation stores this pointer of p */
	nnsliteBlistBegin(blist, &iter);
	while(nnsliteBlistNext(&iter, &record) == NNSLITE_RET_TRUE)
	{
		if (record.addr == p)
		{
			nnsliteBlistRemove(&iter);
			gSystemMemContext.freeMem(p);
			return;
		}
	}
	gSystemMemContext.freeMem(p);
	printf("This memory block %x is not recorded!\n", (int)p);
	assert(false);
}

static NNSLITE_RET DefaultCheckMem()
{
	u32 non_freed = 0;
	size_t leak_memsize = 0;
	AllocRecord record;
	nnsliteBlistIterator iter;
	nnsliteBlist* blist = &gDefaultMemContext.allocList;
	nnsliteBlistBegin(blist, &iter);
	while(nnsliteBlistNext(&iter, &record) == NNSLITE_RET_TRUE)
	{
		printf("This memory block is not released : %lld\n", record.id);
		non_freed++;
		leak_memsize += record.size;
	}
	printf("Totally %d memory blocks, %d bytes are leaked\n", non_freed, leak_memsize);
	return NNSLITE_RET_OK;
}


